/*
 * Copyright (c) 1989 - 1999 Microsoft Corporation
 */

#include <windows.h>
#include <windef.h>
#include <winbase.h>
#include <winsvc.h>
#include <winnetwk.h>
#include <npapi.h>
#include <devioctl.h>

#include <strsafe.h>
#include <list>
#include <string>
#include <algorithm>
#include <vector>
#include <array>
#include <cassert>
#include "../hwfs/hwfs_shared.h"
#include "hwfs_enum.h"
#include "debug.h"

#pragma warning(disable: 4100)

#define WNNC_DRIVER(major, minor) ((major) * 0x00010000 + (minor))

constexpr const size_t LOCAL_DRIVE_LEN = 2;
constexpr const size_t LOCAL_DRIVE_BUFFER_SIZE = (LOCAL_DRIVE_LEN + 1) * sizeof(wchar_t);
constexpr const size_t LOCAL_DRIVE_SIZE = LOCAL_DRIVE_LEN * sizeof(wchar_t);

using LocalDrive = std::array<wchar_t, LOCAL_DRIVE_LEN + 1>;

extern "C"
DWORD APIENTRY
NPGetCaps(DWORD nIndex)
{
    // return the capaboilities of the Mini redirector
    DWORD ret = 0;

    switch (nIndex) {
        case WNNC_SPEC_VERSION:
            ret = WNNC_SPEC_VERSION51;
            break;

        case WNNC_NET_TYPE:
            ret = WNNC_NET_RDR2SAMPLE | 0x7313; // Actually, need contact Microsoft and get our unique number
            break;

        case WNNC_DRIVER_VERSION:
            ret = WNNC_DRIVER(1, 0);
            break;

        case WNNC_CONNECTION:
            ret = WNNC_CON_GETCONNECTIONS |
                 WNNC_CON_CANCELCONNECTION |
                 WNNC_CON_ADDCONNECTION |
                 WNNC_CON_ADDCONNECTION3;
            break;

        case WNNC_ENUMERATION:
            ret = WNNC_ENUM_LOCAL |
                 WNNC_ENUM_GLOBAL;
            break;

        case WNNC_START:
            ret = 1;
            break;

        case WNNC_DIALOG:
            ret = WNNC_DLG_GETRESOURCEINFORMATION;
            break;

        case WNNC_USER:
        case WNNC_ADMIN:
        default:
            ret = 0;
            break;
    }
    DBG_PRINT((L"GetNetCaps[%u]: 0x%08x\n", nIndex, ret));

    return ret;
}


extern "C"
DWORD APIENTRY
NPLogonNotify(
    IN  PLUID   lpLogonId,
    IN  PCWSTR lpAuthentInfoType,
    IN  PVOID  lpAuthentInfo,
    IN  PCWSTR lpPreviousAuthentInfoType,
    IN  PVOID  lpPreviousAuthentInfo,
    IN  PWSTR  lpStationName,
    IN  PVOID  StationHandle,
    OUT PWSTR  *lpLogonScript)
{
    // Handle logon notification
    *lpLogonScript = NULL;

    return WN_SUCCESS;
}


extern "C"
DWORD APIENTRY
NPPasswordChangeNotify(
    IN  LPCWSTR lpAuthentInfoType,
    IN  LPVOID  lpAuthentInfo,
    IN  LPCWSTR lpPreviousAuthentInfoType,
    IN  LPVOID  lpPreviousAuthentInfo,
    IN  LPWSTR  lpStationName,
    LPVOID  StationHandle,
    DWORD   dwChangeInfo)
{
    // Handle password change notification
    SetLastError(WN_NOT_SUPPORTED);

    return WN_NOT_SUPPORTED;
}


ULONG SendToMiniRdr(
    IN ULONG            IoctlCode,
    IN PVOID            InputDataBuf,
    IN ULONG            InputDataLen,
    IN PVOID            OutputDataBuf,
    IN PULONG           pOutputDataLen)
{
    // Send device ioctl to kernel driver
    HANDLE  DeviceHandle;
    ULONG   Status;

    Status = WN_SUCCESS;

    DeviceHandle = CreateFile(
        DD_HWFS_USERMODE_DEV_NAME_U, GENERIC_READ | GENERIC_WRITE,
        FILE_SHARE_READ | FILE_SHARE_WRITE, (LPSECURITY_ATTRIBUTES)NULL,
        OPEN_EXISTING, 0, (HANDLE)NULL);
    if (INVALID_HANDLE_VALUE != DeviceHandle) {
        BOOL rc = DeviceIoControl(
            DeviceHandle, IoctlCode,
            InputDataBuf, InputDataLen,
            OutputDataBuf, *pOutputDataLen,
            pOutputDataLen, NULL);
        if (!rc) {
            Status = GetLastError();
            DBG_PRINT((L"SendToMiniRdr: DeviceIoctl error: %u\n", Status));
        } else {
            DBG_PRINT((L"SendToMiniRdr: DeviceIoctl success: BytesReturned %u\n", *pOutputDataLen));
        }
        CloseHandle(DeviceHandle);
    } else {
        Status = GetLastError();
        DBG_PRINT((L"SendToMiniRdr: error 0x%lx opening device\n", Status));
    }

    return Status;
}


extern "C"
DWORD APIENTRY
NPAddConnection(
    IN  LPNETRESOURCE   lpNetResource,
    IN  LPWSTR          lpPassword OPTIONAL,
    IN  LPWSTR          lpUserName OPTIONAL)
{
    // Add a connection to the list of connections
    // associated with this network provider

    DBG_PRINT((L"NPAddConnection....\n"));

    return NPAddConnection3(NULL, lpNetResource, lpPassword, lpUserName, 0);
}

static const unsigned NP_NAME_SIZE = 128;
static const unsigned DRIVE_COLON_POS = 1;

static bool HwfsLocalNameString(const WCHAR *wszLocalName, std::wstring &localName)
{
    localName.clear();
    if (wszLocalName && wcslen(wszLocalName) > 1) {
        if (wszLocalName[DRIVE_COLON_POS] != L':') {
            return false;
        }
        localName = L"C:";

        localName[0] = (WCHAR)toupper(wszLocalName[0]);
    }
    return true;
}

static const WCHAR SLASHSLASH_STRING[] = L"\\\\";
static const unsigned SLASHSLASH_LEN = 2;

static bool HwfsRemoteNameString(const WCHAR* wszRemoteName, std::wstring& remoteName)
{
    remoteName.clear();
    if (!wszRemoteName || wcslen(wszRemoteName) < SLASHSLASH_LEN ||
            wszRemoteName[0] != L'\\' || wszRemoteName[1] != L'\\') {
        return false;
    }
    remoteName = (wszRemoteName + SLASHSLASH_LEN);
    return true;
}

// Create name in format \device\miniredirector\;<DriveLetter>:\Server\Share
// or \device\miniredirector\;\Server
static DWORD HwfsConnectionNameString(
    const WCHAR* wszRemoteName, const WCHAR* wszLocalName, std::wstring& connectionName)
{
    std::wstring localName;
    bool bRes = HwfsLocalNameString(wszLocalName, localName);
    if (!bRes) {
        return WN_BAD_LOCALNAME;
    }

    std::wstring remoteName;
    bRes = HwfsRemoteNameString(wszRemoteName, remoteName);
    if (!bRes) {
        return WN_BAD_NETNAME;
    }

    connectionName = DD_HWFS_FS_DEVICE_NAME_U + std::wstring(L"\\;");
    connectionName += localName;
    connectionName += std::wstring(L"\\") + remoteName;

    return WN_SUCCESS;
}

extern "C"
DWORD APIENTRY
NPAddConnection3(
    IN  HWND            hwndOwner,
    IN  LPNETRESOURCE   lpNetResource,
    IN  LPWSTR          lpPassword OPTIONAL,
    IN  LPWSTR          lpUserName OPTIONAL,
    IN  DWORD           dwFlags)
{
    // Add a connection to the list of connections
    // associated with this network provider
    DBG_PRINT((L"NPAddConnection3....\n"));

    DBG_PRINT((L"Local Name:  %s\n", lpNetResource->lpLocalName));
    DBG_PRINT((L"Remote Name: %s\n", lpNetResource->lpRemoteName));
    DBG_PRINT((L"ResourceType: %x\n", lpNetResource->dwType));

    if (lpNetResource->dwType != RESOURCETYPE_DISK && lpNetResource->dwType != RESOURCETYPE_ANY) {
        return WN_BAD_NETNAME;
    }

    // Create name in format \device\miniredirector\;<DriveLetter>:\Server\Share
    // or \device\miniredirector\;\Server
    std::wstring connectionName;
    DWORD status = HwfsConnectionNameString(
        lpNetResource->lpRemoteName, lpNetResource->lpLocalName, connectionName);
    if (status != WN_SUCCESS) {
        return status;
    }

    DBG_PRINT((L"Full Connect Name: %s; StringLength %d\n",
        connectionName.c_str(), (connectionName.length() + 1) * sizeof(WCHAR)));

    SetLastError(NO_ERROR);

    WCHAR wszScratch[NP_NAME_SIZE];
    std::wstring localName;
    bool bRes = HwfsLocalNameString(lpNetResource->lpLocalName, localName);
    if (!bRes) {
        return WN_BAD_NETNAME;
    }

    if (!localName.empty() &&
            QueryDosDevice(localName.c_str(), wszScratch, NP_NAME_SIZE)) {
        return WN_ALREADY_CONNECTED;
    }

    status = WN_SUCCESS;
    if (localName.empty() || GetLastError() == ERROR_FILE_NOT_FOUND) {
        DWORD copyBytes = 0;
        status = SendToMiniRdr(IOCTL_HWFS_ADDCONN,
                    (PVOID)connectionName.c_str(),
                    (ULONG)((connectionName.length() + 1) * sizeof(WCHAR)),
                    NULL, &copyBytes);
        if (status != WN_SUCCESS) {
            return WN_BAD_NETNAME;
        }

        if (!localName.empty()) {
            if (!DefineDosDevice(DDD_RAW_TARGET_PATH |
                    DDD_NO_BROADCAST_SYSTEM,
                    localName.c_str(),
                    connectionName.c_str())) {
                status = GetLastError();
            }
        }
    }

    DBG_PRINT((L"NPAddConnection3 finished: %u\n", status));

    return status;
}

std::wstring GetCurrentProcessName()
{
    ULONG currSize = MAX_PATH + 1;
    std::vector<WCHAR> moduleName(currSize);
    while (true) {
        ULONG retSize = GetModuleFileName(NULL, moduleName.data(), currSize);
        if (retSize == 0) {
            return L"";
        }
        if (retSize < currSize) {
            return moduleName.data();
        }
        currSize += MAX_PATH;
        moduleName.resize(currSize);
    }
}

extern "C"
DWORD APIENTRY
NPCancelConnection(
    IN  LPWSTR  lpName,
    IN  BOOL    fForce)
{
    // Delete a connection from the list of connections
    // associated with this network provider
    std::wstring processName = GetCurrentProcessName();

    DBG_PRINT((L"NPCancelConnection: %s from process %s; force %d\n",
        lpName, processName.c_str(), (int)fForce));

    std::size_t slashPos = processName.find_last_of(L"/\\");
    processName = processName.substr(slashPos + 1);
    if (0 == _wcsicmp(processName.c_str(), L"winlogon.exe")) {
        // Ignore requests from winlogon
        return WN_NOT_CONNECTED;
    }

    std::wstring localName;
    bool bRes = HwfsLocalNameString(lpName, localName);
    if (!bRes || localName.empty()) {
        return WN_NOT_CONNECTED;
    }

    WCHAR wszRemoteName[NP_NAME_SIZE];
    ULONG copyBytes = NP_NAME_SIZE * sizeof(WCHAR);

    // GETCONN returns size with trailing zero.
    DWORD status = SendToMiniRdr(IOCTL_HWFS_GETCONN,
        (PVOID)localName.c_str(), (ULONG)(localName.size() + 1) * sizeof(WCHAR),
        (PVOID)wszRemoteName, &copyBytes);
    if (status != WN_SUCCESS) {
        return WN_NOT_CONNECTED;
    }
    if (copyBytes == 0 || copyBytes > NP_NAME_SIZE * sizeof(WCHAR)) {
        return WN_NET_ERROR;
    }
    wszRemoteName[copyBytes / sizeof(WCHAR) - 1] = L'\0';
    assert(wszRemoteName[0] == L'\\');

    std::wstring remoteName = wszRemoteName + 1;

    if (!DefineDosDeviceW(DDD_REMOVE_DEFINITION | DDD_RAW_TARGET_PATH,
            localName.c_str(), NULL)) {
        status = GetLastError();
        DBG_PRINT((L"NPCancelConnection Remove DosDevice error %u\n", status));
        return ERROR_ACCESS_DENIED;
    }

    std::wstring connectionName = DD_HWFS_FS_DEVICE_NAME_U + std::wstring(L"\\;");
    connectionName += localName;
    connectionName += std::wstring(L"\\") + remoteName;

    copyBytes = 0;
    status = SendToMiniRdr(IOCTL_HWFS_DELCONN, (PVOID)connectionName.c_str(),
                (ULONG)(connectionName.size() + 1) * sizeof(WCHAR),
                NULL, &copyBytes);
    if (status != WN_SUCCESS) {
        status = WN_NOT_CONNECTED;
    }

    DBG_PRINT((L"NPCancelConnection finished: %u\n", status));
    return status;
}


extern "C"
DWORD APIENTRY
NPGetConnection(
    IN  LPWSTR  lpLocalName,
    OUT LPWSTR  lpRemoteName,
    IN OUT LPDWORD lpBufferSize)
{
    // return information associated with a connection
    DBG_PRINT((L"NPGetConnection for '%ws'....\n", lpLocalName));

    if (lpLocalName == NULL) {
        return WN_NOT_CONNECTED;
    }

    std::wstring localName;
    HwfsLocalNameString(lpLocalName, localName);
    if (localName.empty()) {
        return WN_NOT_CONNECTED;
    }

    WCHAR wszRemoteName[NP_NAME_SIZE];
    ULONG copyBytes = NP_NAME_SIZE * sizeof(WCHAR);
    PVOID inputDataBuf = (PVOID)localName.c_str();
    ULONG inputDataLen = (ULONG)(localName.size() + 1) * sizeof(WCHAR);

    DWORD status = SendToMiniRdr(IOCTL_HWFS_GETCONN, inputDataBuf, inputDataLen, wszRemoteName, &copyBytes);
    if (status != NO_ERROR) {
	    return copyBytes == 0 ? WN_NO_NETWORK : WN_NOT_CONNECTED;
	}

    assert(wszRemoteName[0] == L'\\');
	wszRemoteName[copyBytes / sizeof(WCHAR)] = 0;

    // The wszRemoteName is like "\SRV". We need return "\\SRV".
    std::wstring remoteName = std::wstring(SLASHSLASH_STRING).append(wszRemoteName + 1);
    size_t len = (remoteName.size() + 1) * sizeof(WCHAR);
    if (*lpBufferSize < len) {
         *lpBufferSize = (ULONG)len;
         return WN_MORE_DATA;
    }

    errno_t err = wcscpy_s(lpRemoteName, *lpBufferSize / sizeof(WCHAR), remoteName.c_str());
    if (err) {
        DBG_PRINT((L"NPGetConnection: internal error\n"));
        return ERROR_INTERNAL_ERROR;
    }

    lpRemoteName[remoteName.size()] = L'\0';
    DBG_PRINT((L"NPGetConnection: '%s' -> '%s'\n", localName.c_str(), lpRemoteName));

    return WN_SUCCESS;
}


extern "C"
DWORD APIENTRY
NPOpenEnum(
    IN  DWORD          dwScope,
    IN  DWORD          dwType,
    IN  DWORD          dwUsage,
    IN  LPNETRESOURCE  lpNetResource,
    OUT LPHANDLE       lphEnum)
{
    // This routine opens a handle for enumeration of resources.
    // We intentionally enumerate only connected shares.
    DBG_PRINT((L"NPOpenEnum\n"));

    *lphEnum = NULL;

    NPEnum* en = new NPEnum(dwScope, dwType, dwUsage, lpNetResource);
    DWORD status = en->EnumerateResources();
    if (status) {
        delete en;
        return status;
    }

    *lphEnum = en;
    return WN_SUCCESS;
}


extern "C"
DWORD APIENTRY
NPEnumResource(
    IN     HANDLE  hEnum,
    IN OUT LPDWORD lpcCount,
    OUT    LPVOID  lpBuffer,
    IN OUT LPDWORD lpBufferSize)
{
    // Return next portion of resources
    NPEnum* en = (NPEnum *)hEnum;
    if (!en) {
        return WN_BAD_HANDLE;
    }
    if (!lpcCount || !lpBuffer || !lpBufferSize) {
        return WN_BAD_VALUE;
    }

    DWORD Status = en->FillResourceRequest(lpcCount, lpBuffer, lpBufferSize);

    DBG_PRINT((L"NPEnumResource Entries returned - %d entries, status %d\n",
        *lpcCount, Status));

    return Status;
}


extern "C"
DWORD APIENTRY
NPCloseEnum(
    IN  HANDLE hEnum)
{
    // Close enumeration handle
    DBG_PRINT((L"NPCloseEnum\n"));

    NPEnum* en = (NPEnum*)hEnum;
    delete en;

    return WN_SUCCESS;
}


extern "C"
DWORD APIENTRY
NPGetResourceParent(
    LPNETRESOURCE   lpNetResource,
    LPVOID  lpBuffer,
    LPDWORD lpBufferSize)
{
    // Return information about resource parent
    DBG_PRINT((L"NPGetResourceParent: WN_NOT_SUPPORTED\n"));

    return WN_NOT_SUPPORTED;
}

static DWORD SplitResourceRemoteName(
    PWCHAR remoteName,
    PWCHAR& srvShareName, ULONG& srvShareNameLen,
    PWCHAR& sysName, ULONG& sysNameLen,
    bool& serverOnlyName)
{
    size_t remoteNameLen = wcsnlen(remoteName, MAX_PATH);
    if (remoteNameLen == MAX_PATH) {
        return WN_BAD_VALUE;
    }

    const DWORD hwfsRootSrvLen = (DWORD)wcslen(HWFS_ROOTSRV_NAME);
    if (0 != _wcsnicmp(remoteName, HWFS_ROOTSRV_NAME, hwfsRootSrvLen)) {
        return WN_BAD_NETNAME;
    }

    sysName = NULL;
    sysNameLen = 0;

    srvShareName = remoteName;
    srvShareNameLen = hwfsRootSrvLen;

    serverOnlyName = true;

    // Name must be either "\\hwfs" or "\\hwfs\", or be like "\\hwfs\Share + {some path}"
    PWCHAR afterSrvPath = remoteName + hwfsRootSrvLen;
    if (afterSrvPath[0] == 0) {
        return WN_SUCCESS;
    }
    if (afterSrvPath[0] != L'\\') {
        return WN_BAD_NETNAME;
    }
    afterSrvPath++;
    if (afterSrvPath[0] == 0) {
        return WN_SUCCESS;
    }

    serverOnlyName = false;

    // nextSlash points to the string after "\\hwfs\\Share"
    sysName = wcschr(afterSrvPath, L'\\');
    if (sysName == NULL) {
        // It's a share name without sys-path. It's ok.
        srvShareNameLen = (ULONG)remoteNameLen;
        return WN_SUCCESS;
    }

    sysNameLen = (ULONG)(remoteNameLen - (sysName - remoteName));
    return WN_SUCCESS;
}

static PWCHAR CopyStringToBuf(PWCHAR& dst, size_t& dstMaxSizeChars, PWCHAR src, size_t srcSizeChars)
{
    if (dstMaxSizeChars < srcSizeChars + 1) {
        DBG_PRINT((L"No size: %u < %u\n", dstMaxSizeChars, srcSizeChars + 1));
        return NULL;
    }

    errno_t err = wmemcpy_s(dst, dstMaxSizeChars, src, srcSizeChars);
    if (err != 0) {
        DBG_PRINT((L"wmemcpy_s failed: %u : %u\n", dstMaxSizeChars, srcSizeChars + 1));
        return NULL;
    }
    PWCHAR ret = dst;

    dst[srcSizeChars] = L'\0';
    dst += srcSizeChars + 1; // include trailing zero
    dstMaxSizeChars -= srcSizeChars + 1;
    return ret;
}

extern "C"
DWORD APIENTRY
NPGetResourceInformation(
    IN     LPNETRESOURCE   lpNetResource,
    OUT    LPVOID          lpBuffer,
    IN OUT LPDWORD         lpBufferSize,
    OUT    LPWSTR*         lplpSystem)
{
    // Return information associated with resource
    bool bServerOnlyName;
    PWCHAR wszServerShare;
    ULONG  serverShareLen;

    PWCHAR wszSysPath;
    ULONG sysPathLen;

    // We must split the string in pNetResourceInfo->lpRemoteName in "//Server/Share" and "/Directory" path.
    // For example, if name is "//HwFs/Share/Dir1/Dir2", lpRemoteName must point to "//HwFs/Share", and
    // and lplpSystem to "/Dir1/Dir2".
    DWORD status = SplitResourceRemoteName(lpNetResource->lpRemoteName,
        wszServerShare, serverShareLen,
        wszSysPath, sysPathLen,
        bServerOnlyName);
    if (status != WN_SUCCESS) {
        DBG_PRINT((L"NPGetResourceInformation: Failed to split remote name %s\n", lpNetResource->lpRemoteName));
        return status;
    }

    DBG_PRINT((L"RemoteName is %s\n", lpNetResource->lpRemoteName));

    ULONG hwfsProviderNameLen = (ULONG)wcslen(HWFS_PROVIDER_NAME_U);

    ULONG allStringsLen = hwfsProviderNameLen + serverShareLen + sysPathLen + 1 + 1 + 1; // 3 trailing 0

    ULONG fullSize = (ULONG)sizeof(NETRESOURCE) + allStringsLen * sizeof(WCHAR);
    if (*lpBufferSize < fullSize) {
        DBG_PRINT((L"WN_MORE_DATA: got %u, need %u\n", *lpBufferSize, fullSize));
        *lpBufferSize = fullSize;
        return WN_MORE_DATA;
    }

    LPNETRESOURCE pNetResourceInfo = (LPNETRESOURCE)lpBuffer;
    memset(pNetResourceInfo, 0, sizeof(NETRESOURCE));

    rsize_t lengthRemain = allStringsLen;
    WCHAR* pStrPtr = (WCHAR*)(pNetResourceInfo + 1);

    pNetResourceInfo->lpRemoteName = CopyStringToBuf(pStrPtr, lengthRemain, wszServerShare, serverShareLen);
    if (!pNetResourceInfo->lpRemoteName) {
        return WN_BAD_VALUE;
    }

    pNetResourceInfo->lpProvider = CopyStringToBuf(pStrPtr, lengthRemain, HWFS_PROVIDER_NAME_U, hwfsProviderNameLen);
    if (!pNetResourceInfo->lpProvider) {
        return WN_BAD_VALUE;
    }

    if (bServerOnlyName) {
        pNetResourceInfo->dwType = RESOURCETYPE_ANY;
        pNetResourceInfo->dwDisplayType = RESOURCEDISPLAYTYPE_SERVER;
        pNetResourceInfo->dwUsage = RESOURCEUSAGE_CONTAINER;
        if (lplpSystem) {
            *lplpSystem = NULL;
        }
        DBG_PRINT((L"Server-only request: Srv = %s\n", pNetResourceInfo->lpRemoteName));
        return WN_SUCCESS;
    }

    pNetResourceInfo->dwType = RESOURCETYPE_DISK;
    pNetResourceInfo->dwDisplayType = RESOURCEDISPLAYTYPE_SHARE;
    pNetResourceInfo->dwUsage = RESOURCEUSAGE_CONNECTABLE;

    if (!lplpSystem) {
        DBG_PRINT((L"No sys-path component requested: Share = %s\n", pNetResourceInfo->lpRemoteName));
        return WN_SUCCESS;
    }
    if (sysPathLen == 0) {
        DBG_PRINT((L"sys-path component is zero: Share = %s\n", pNetResourceInfo->lpRemoteName));
        lplpSystem = NULL;
        return WN_SUCCESS;
    }
    PWCHAR sysPtr = CopyStringToBuf(pStrPtr, lengthRemain, wszSysPath, sysPathLen);
    if (!sysPtr) {
        return WN_BAD_VALUE;
    }

    *lplpSystem = sysPtr;
    DBG_PRINT((L"Share = '%s', SysPath = '%s'\n", pNetResourceInfo->lpRemoteName, sysPtr));

    return WN_SUCCESS;
}

static LocalDrive GetLocalDrive(
    IN     LPCWSTR pLocalPath)
{
    LocalDrive drive;
    drive.fill(0);
    memcpy_s(drive.data(), LOCAL_DRIVE_BUFFER_SIZE, pLocalPath, LOCAL_DRIVE_SIZE);
    return drive;
}

static DWORD GetUniversalName(
    IN     LPCWSTR pLocalPath,
    OUT    LPVOID  pBuffer,
    IN OUT LPDWORD pBufferSize)
{
    assert(pLocalPath != NULL);
    auto localDrive = GetLocalDrive(pLocalPath);

    LPCWSTR pRemainingPath = &pLocalPath[2];
    DWORD remainingPathLength = (DWORD)wcslen(pRemainingPath);
    DWORD remainingPathSize = remainingPathLength * sizeof(WCHAR);

    LPUNIVERSAL_NAME_INFO pUniversalNameInfo = (LPUNIVERSAL_NAME_INFO)pBuffer;
    pUniversalNameInfo->lpUniversalName = (PWCHAR)(pUniversalNameInfo + 1);

    DWORD remoteNameBufferSize = 0;
    if (*pBufferSize > sizeof(UNIVERSAL_NAME_INFO)) {
        remoteNameBufferSize = *pBufferSize - sizeof(UNIVERSAL_NAME_INFO);
    }

    DWORD dwStatus = NPGetConnection(localDrive.data(), pUniversalNameInfo->lpUniversalName, &remoteNameBufferSize);
    if (dwStatus != WN_SUCCESS) {
        if (dwStatus == WN_MORE_DATA) {
            DBG_PRINT((L"NPGetUniversalName: buffer is too small.\n"));
            *pBufferSize = remoteNameBufferSize + remainingPathSize + sizeof(UNIVERSAL_NAME_INFO);
            return WN_MORE_DATA;
        }
        return dwStatus;
    }

    DWORD remoteNameLength = (DWORD)wcslen(pUniversalNameInfo->lpUniversalName);
    PWCHAR pString = &pUniversalNameInfo->lpUniversalName[remoteNameLength];

    DWORD requiredBufferSize = remoteNameLength * sizeof(WCHAR) + remainingPathSize + sizeof(WCHAR);
    if (*pBufferSize < requiredBufferSize) {
        *pBufferSize = requiredBufferSize;
        return WN_MORE_DATA;
    }
    memcpy_s(pString, remoteNameBufferSize, pRemainingPath, remainingPathSize);
    pUniversalNameInfo->lpUniversalName[remoteNameLength + remainingPathLength] = 0;
    return WN_SUCCESS;
}

static DWORD CalculateRemoteNameBufferSize(DWORD connectionNameSize, DWORD remainingPathSize)
{
    DWORD structSize = sizeof(REMOTE_NAME_INFO);
    DWORD universalNameBufferSize = connectionNameSize + remainingPathSize + sizeof(WCHAR);
    DWORD connectionNameBufferSize = connectionNameSize + sizeof(WCHAR);
    DWORD remainingPathBufferSize = remainingPathSize + sizeof(WCHAR);

    return structSize + universalNameBufferSize + connectionNameBufferSize + remainingPathBufferSize;
}

static DWORD GetRemoteName(
    IN     LPCWSTR pLocalPath,
    OUT    LPVOID  pBuffer,
    IN OUT LPDWORD pBufferSize)
{
    assert(pLocalPath != NULL);
    auto localDrive = GetLocalDrive(pLocalPath);

    LPCWSTR pRemainingPath = &pLocalPath[2];
    DWORD remainingPathLength = (DWORD)(wcslen(pRemainingPath));
    DWORD remainingPathSize = remainingPathLength * sizeof(WCHAR);

    DWORD remoteNameBufferSize = 0;
    LPWSTR lpUniversalName = NULL;
    LPREMOTE_NAME_INFO pRemoteNameInfo = (LPREMOTE_NAME_INFO)pBuffer;
    if (*pBufferSize > sizeof(REMOTE_NAME_INFO)) {
        lpUniversalName = (PWCHAR)(pRemoteNameInfo + 1);
        pRemoteNameInfo->lpUniversalName = lpUniversalName;

        remoteNameBufferSize = *pBufferSize - sizeof(REMOTE_NAME_INFO);
    }

    DWORD dwStatus = NPGetConnection(localDrive.data(), lpUniversalName, &remoteNameBufferSize);
    if (dwStatus != WN_SUCCESS) {
        if (dwStatus == WN_MORE_DATA) {
            DBG_PRINT((L"NPGetUniversalName: buffer is too small.\n"));
            *pBufferSize = CalculateRemoteNameBufferSize(remoteNameBufferSize - sizeof(WCHAR), remainingPathSize);
            return WN_MORE_DATA;
        }
        return dwStatus;
    }

    DWORD remoteNameLength = (DWORD)wcslen(pRemoteNameInfo->lpUniversalName);
    DWORD remoteNameSize = remoteNameLength * sizeof(WCHAR);

    DWORD requiredBufferSize = CalculateRemoteNameBufferSize(remoteNameSize, remainingPathSize);
    if (*pBufferSize < requiredBufferSize) {
        *pBufferSize = requiredBufferSize;
        return WN_MORE_DATA;
    }

    PWCHAR pString = &pRemoteNameInfo->lpUniversalName[remoteNameLength];
    long spaceInBuffer = remoteNameBufferSize;
    memcpy_s(pString, spaceInBuffer, pRemainingPath, remainingPathSize);

    spaceInBuffer -= remainingPathSize;
    pString += remainingPathLength;
    *pString++ = 0;

    pRemoteNameInfo->lpConnectionName = pString;
    memcpy_s(pString, spaceInBuffer, pRemoteNameInfo->lpUniversalName, remoteNameSize);

    spaceInBuffer -= remoteNameSize;
    pString += remoteNameLength;
    *pString++ = 0;

    pRemoteNameInfo->lpRemainingPath = pString;
    memcpy_s(pString, spaceInBuffer, pRemainingPath, remainingPathSize);

    return WN_SUCCESS;
}

extern "C"
DWORD APIENTRY
NPGetUniversalName(
    IN     LPCWSTR pLocalPath,
    IN     DWORD   dwInfoLevel,
    OUT    LPVOID  pBuffer,
    IN OUT LPDWORD pBufferSize)
{
    if (pLocalPath == NULL || pLocalPath[0] == 0 || pLocalPath[1] != L':') {
        DBG_PRINT((L"NPGetUniversalName: Bad pLocalPath.\n"));
        return WN_BAD_LOCALNAME;
    }

    DBG_PRINT((L"NPGetUniversalName: pLocalPath = %ls, InfoLevel = %d, *pBufferSize = %d\n",
         pLocalPath, dwInfoLevel, *pBufferSize));

    switch (dwInfoLevel) {
        case UNIVERSAL_NAME_INFO_LEVEL:
            return GetUniversalName(pLocalPath, pBuffer, pBufferSize);
        case REMOTE_NAME_INFO_LEVEL:
            return GetRemoteName(pLocalPath, pBuffer, pBufferSize);
        default:
            DBG_PRINT((L"NPGetUniversalName: Bad dwInfoLevel value: %d\n", dwInfoLevel));
            return WN_BAD_LEVEL;
    }
}

static const UINT MAX_DEBUG_STRING_CHARS = 256;

// Format and write debug information to OutputDebugString
ULONG _cdecl WideDbgPrint(IN  LPTSTR Format, ...)
{
    ULONG rc = 0;
    WCHAR szbuffer[MAX_DEBUG_STRING_CHARS];

    va_list marker;
    va_start(marker, Format);
    {
        StringCchVPrintfW(szbuffer, MAX_DEBUG_STRING_CHARS - 1, Format, marker);
        szbuffer[MAX_DEBUG_STRING_CHARS - 1] = (TCHAR)0;
        OutputDebugString(TRACE_TAG);
        OutputDebugString(szbuffer);
    }

    return rc;
}
